import backtrader as bt
import pandas as pd
import numpy as np
from datetime import datetime
import os

# 创建自定义数据源类
class TushareCSVData(bt.feeds.GenericCSVData):
    """自定义数据源类，用于读取tushare格式的CSV文件"""
    
    params = (
        ('dtformat', '%Y%m%d'),  # 日期格式
        ('datetime', 1),          # 日期列索引
        ('open', 2),              # 开盘价列索引
        ('high', 3),              # 最高价列索引
        ('low', 4),               # 最低价列索引
        ('close', 5),             # 收盘价列索引
        ('volume', 9),            # 成交量列索引
        ('openinterest', -1),     # 未使用
        ('headers', True),        # 有表头
        ('encoding', 'utf-8'),    # UTF-8编码
        ('separator', ','),       # 分隔符
        ('quotes', True),         # 处理引号
    )
    
    def _loadline(self, linetokens):
        """重写_loadline方法来处理引号"""
        # 去除每个字段的引号
        cleaned_tokens = [token.strip('"') for token in linetokens]
        return super()._loadline(cleaned_tokens)

# 自定义布林带指标类
class CustomBollingerBands(bt.Indicator):
    """按照指定公式计算布林带参数"""
    
    lines = ('mida', 'uppera', 'lowera', 'boll', 'ub', 'lb', 'bb', 'mb')
    params = (('period', 20), ('bb_period', 6))
    
    def __init__(self):
        # 存储历史数据用于计算
        self.close_prices = []
        self.vart1_values = []
        self.li_jin_values = []
        self.mida_values = []
        self.uppera_values = []
        self.lowera_values = []
        self.bb_values = []
    
    def next(self):
        """计算当前值"""
        # 添加当前收盘价
        self.close_prices.append(self.data.close[0])
        
        # 保持历史数据长度
        if len(self.close_prices) > self.params.period * 2:
            self.close_prices = self.close_prices[-self.params.period * 2:]
        
        # 计算MIDA (20日移动平均)
        if len(self.close_prices) >= self.params.period:
            mida_value = sum(self.close_prices[-self.params.period:]) / self.params.period
        else:
            mida_value = self.data.close[0]
        
        self.mida_values.append(mida_value)
        
        # 计算VART1 (收盘价与移动平均的差值的平方)
        vart1_value = (self.data.close[0] - mida_value) ** 2
        self.vart1_values.append(vart1_value)
        
        # 计算李津 (VART1的20日移动平均)
        if len(self.vart1_values) >= self.params.period:
            li_jin_value = sum(self.vart1_values[-self.params.period:]) / self.params.period
        else:
            li_jin_value = vart1_value
        
        self.li_jin_values.append(li_jin_value)
        
        # 计算VART3 (标准差)
        vart3_value = np.sqrt(li_jin_value )
        
        # 计算UPPERA和LOWERA
        uppera_value = mida_value + 2 * vart3_value
        lowera_value = mida_value - 2 * vart3_value
        
        self.uppera_values.append(uppera_value)
        self.lowera_values.append(lowera_value)
        
        # 计算BOLL, UB, LB (前一日值)
        boll_value = self.mida_values[-2] if len(self.mida_values) > 1 else mida_value
        ub_value = self.uppera_values[-2] if len(self.uppera_values) > 1 else uppera_value
        lb_value = self.lowera_values[-2] if len(self.lowera_values) > 1 else lowera_value
        
        # 计算BB值
        if ub_value != lb_value:
            bb_value = (self.data.close[0] - ub_value) / (ub_value - lb_value) * 100 + 75
        else:
            bb_value = 75
        
        self.bb_values.append(bb_value)
        
        # 计算MB值 (BB的6日移动平均)
        if len(self.bb_values) >= self.params.bb_period:
            mb_value = sum(self.bb_values[-self.params.bb_period:]) / self.params.bb_period
        else:
            mb_value = bb_value
        
        # 将计算结果赋值给lines
        self.lines.mida[0] = mida_value
        self.lines.uppera[0] = uppera_value
        self.lines.lowera[0] = lowera_value
        self.lines.boll[0] = boll_value
        self.lines.ub[0] = ub_value
        self.lines.lb[0] = lb_value
        self.lines.bb[0] = bb_value
        self.lines.mb[0] = mb_value

# 创建策略类
class BollingerStrategy(bt.Strategy):
    """布林带策略类"""
    
    def __init__(self):
        # 添加自定义布林带指标
        self.bollinger = CustomBollingerBands()
        
    def next(self):
        """每个交易日调用一次"""
        # 获取当前日期
        current_date = self.data.datetime.date(0)
        
        # 获取布林带参数
        mida = self.bollinger.lines.mida[0]
        uppera = self.bollinger.lines.uppera[0]
        lowera = self.bollinger.lines.lowera[0]
        boll = self.bollinger.lines.boll[0]
        ub = self.bollinger.lines.ub[0]
        lb = self.bollinger.lines.lb[0]
        bb = self.bollinger.lines.bb[0]
        mb = self.bollinger.lines.mb[0]
        close_price = self.data.close[0]
        
        # 打印每日布林带参数
        print(f"日期: {current_date}")
        print(f"  收盘价: {close_price:.2f}")
        print(f"  MIDA (20日移动平均): {mida:.2f}")
        print(f"  UPPERA (上轨): {uppera:.2f}")
        print(f"  LOWERA (下轨): {lowera:.2f}")
        print(f"  BOLL (前日中轨): {boll:.2f}")
        print(f"  UB (前日上轨): {ub:.2f}")
        print(f"  LB (前日下轨): {lb:.2f}")
        print(f"  BB (百分比B): {bb:.2f}")
        print(f"  MB (BB的6日平均): {mb:.2f}")
        print("-" * 50)

def main():
    """主函数"""
    # 创建Cerebro引擎
    cerebro = bt.Cerebro()
    
    # 设置初始资金
    cerebro.broker.setcash(100000.0)
    
    # 获取CSV文件路径
    csv_file = os.path.join('data', 'test.csv')
    
    # 检查文件是否存在
    if not os.path.exists(csv_file):
        print(f"错误: 找不到文件 {csv_file}")
        return
    
    # 读取CSV数据并排序
    print("正在读取和排序数据...")
    df = pd.read_csv(csv_file, encoding='utf-8')
    
    # 确保日期列存在
    if len(df.columns) > 1:  # 假设日期列是第2列（索引1）
        date_col = df.columns[1]  # 获取第2列的名称
        df[date_col] = pd.to_datetime(df[date_col], format='%Y%m%d')
        # 按日期排序，老的日期在前
        df = df.sort_values(by=date_col, ascending=True)
        print(f"数据已按日期排序，共 {len(df)} 条记录")
        print(f"日期范围: {df[date_col].min()} 到 {df[date_col].max()}")
    else:
        print("警告: 无法识别日期列，使用原始数据")
    
    # 创建数据源
    if len(df.columns) > 1:
        # 使用排序后的数据创建pandas数据源
        data = bt.feeds.PandasData(
            dataname=df,
            datetime=date_col,  # 使用排序后的日期列
            open=df.columns[2] if len(df.columns) > 2 else None,  # 开盘价列
            high=df.columns[3] if len(df.columns) > 3 else None,  # 最高价列
            low=df.columns[4] if len(df.columns) > 4 else None,   # 最低价列
            close=df.columns[5] if len(df.columns) > 5 else None, # 收盘价列
            volume=df.columns[9] if len(df.columns) > 9 else None, # 成交量列
            openinterest=None,  # 未使用
            fromdate=datetime(2024, 1, 1),  # 设置起始日期
            todate=datetime(2025, 12, 31),   # 设置结束日期
        )
    else:
        # 如果无法识别列，回退到原始CSV数据源
        data = TushareCSVData(
            dataname=csv_file,
            fromdate=datetime(2024, 1, 1),  # 设置起始日期
            todate=datetime(2025, 12, 31),   # 设置结束日期
            reverse=False  # 数据按时间正序排列
        )
    
    # 添加数据到Cerebro
    cerebro.adddata(data)
    
    # 添加策略
    cerebro.addstrategy(BollingerStrategy)
    
    # 添加分析器
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharpe')
    cerebro.addanalyzer(bt.analyzers.Returns, _name='returns')
    
    print("开始回测...")
    print("=" * 60)
    
    # 运行回测
    results = cerebro.run()
    
    # 获取策略实例
    strategy = results[0]
    
    # 打印最终结果
    print("\n" + "=" * 60)
    print("回测完成!")
    print(f"最终资金: {cerebro.broker.getvalue():.2f}")
    
    # 安全地获取分析结果
    returns_analysis = strategy.analyzers.returns.get_analysis()
    sharpe_analysis = strategy.analyzers.sharpe.get_analysis()
    
    if 'rtot' in returns_analysis:
        print(f"总收益率: {returns_analysis['rtot']:.2%}")
    if 'rnorm100' in returns_analysis:
        print(f"年化收益率: {returns_analysis['rnorm100']:.2f}%")
    if 'sharperatio' in sharpe_analysis and sharpe_analysis['sharperatio'] is not None:
        print(f"夏普比率: {sharpe_analysis['sharperatio']:.4f}")
    else:
        print("夏普比率: 无法计算（数据不足）")
    
    # 绘制结果（可选）
    try:
        cerebro.plot(style='candlestick', barup='red', bardown='green')
    except Exception as e:
        print(f"绘图失败: {e}")

if __name__ == '__main__':
    main()
